""""
定义模型
"""
from lib import  ws, max_len
import torch.nn as nn
import torch.nn.functional as F
from  torch.optim import Adam
from dataset import get_dataloder

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(len(ws),100)  # 向量维度100
        self.fc = nn.Linear(max_len*100,2)


    def forward(self,input):
        """

        :param input: [batch_size,max_len]
        :return:
        """
        # 输入必须是longtensor
        x = self.embedding(input) # 进行embedding操作，形状：[batch_size,max_len,100]

        x = x.view([-1, max_len*100])


        out = self.fc(x)
        return F.log_softmax(out,dim=-1)

model = MyModel()
optimizer = Adam(model.parameters(),0.001)

def train(epoch):
    """
    训练
    :param epoch:
    :return:
    """

    for idx,(input,target) in enumerate(get_dataloder()):
        optimizer.zero_grad()
        output = model(input)
        # print(output)
        loss = F.nll_loss(output,target)
        print(loss.item())
        loss.backward()

        optimizer.step()

# def test():





if __name__ =="__main__":

    for i in range(100):
        train(i)